package org.bing.engine.jmx;

import java.io.File;
import java.io.IOException;
import java.util.Properties;

import javax.management.remote.JMXConnector;
import javax.management.remote.JMXConnectorFactory;
import javax.management.remote.JMXServiceURL;

import com.sun.tools.attach.AgentInitializationException;
import com.sun.tools.attach.AgentLoadException;
import com.sun.tools.attach.AttachNotSupportedException;
import com.sun.tools.attach.VirtualMachine;

public class JmxConnectorManager {
    private static final String LOCAL_CONNECTOR_ADDRESS_PROP = "com.sun.management.jmxremote.localConnectorAddress";

    public JMXConnector getJMXConnector(LocalVirtualMachine lvm) {
        try {
            return doGetJMXConnector(lvm);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public JMXConnector doGetJMXConnector(LocalVirtualMachine lvm) throws IOException {
        if (lvm.getAddress() == null) {
            if (lvm.isAttachSupported()) {
                loadManagementAgent(lvm);
            }
        }
        JMXServiceURL jmxUrl = new JMXServiceURL(lvm.getAddress());
        JMXConnector jmxc = JMXConnectorFactory.connect(jmxUrl);
        return jmxc;
    }

    // load the management agent into the target VM
    private void loadManagementAgent(LocalVirtualMachine lvm) throws IOException {
        VirtualMachine vm = null;
        String name = String.valueOf(lvm.getVmid());
        try {
            vm = VirtualMachine.attach(name);
        } catch (AttachNotSupportedException x) {
            throw new IOException(x); // not support jvm, ignore
        }

        String home = vm.getSystemProperties().getProperty("java.home");
        // Normally in ${java.home}/jre/lib/management-agent.jar but might
        // be in ${java.home}/lib in build environments.
        String agent = home + File.separator + "jre" + File.separator + "lib" + File.separator + "management-agent.jar";
        File f = new File(agent);
        if (!f.exists()) {
            agent = home + File.separator + "lib" + File.separator + "management-agent.jar";
            f = new File(agent);
            if (!f.exists()) {
                throw new IOException("Management agent not found");
            }
        }
        agent = f.getCanonicalPath();
        try {
            vm.loadAgent(agent, "com.sun.management.jmxremote");
        } catch (AgentLoadException x) {
            IOException ioe = new IOException(x.getMessage());
            ioe.initCause(x);
            throw ioe;
        } catch (AgentInitializationException x) {
            IOException ioe = new IOException(x.getMessage());
            ioe.initCause(x);
            throw ioe;
        }

        // get the connector address
        Properties agentProps = vm.getAgentProperties();
        lvm.setAddress((String) agentProps.get(LOCAL_CONNECTOR_ADDRESS_PROP));
        vm.detach();
    }
}
